#include "grid.h"

static const int kRandomSeed = 11132;
static const float kStepSize = .20f;
bool Grid::functionTableInit = false;
static bool (*_initializeRuntime[kGridNumSupportedRuntimes])(void);
static bool (*_deinitializeRuntime[kGridNumSupportedRuntimes])(void);
static void (*_computeTotalEnergy[kGridNumSupportedRuntimes])(int,
                                                              float*,
                                                              float*,
                                                              float*,
                                                              int,
                                                              float*,
                                                              float*,
                                                              float*);

//Dummy function used for blank initialization on CPU
static bool GridDoNothing()
{
}

//If not a runtime is not implemented, its function pointer should point to this
static bool GridRuntimeNotImplemented()
{
  throw GridException("This runtime is not supported as compiled.");
}

//Overload of above
static bool GridRuntimeNotImplemented(int numParticles)
{
  throw GridException("This runtime is not supported as compiled.");
}

//Same as above but overloaded for the compute function pointer
static void GridRuntimeNotImplemented(int devNum,
                                      float* x,
                                      float* y,
                                      float* z,
                                      int numParticles,
                                      float* waveFunction,
                                      float* potential,
                                      float* kinetic)
{
  throw GridException("This runtime is not supported as compiled.");
}

void Grid::computeTotalEnergy(bool useCurrentBuffers)
{
  int bufferToUse = useCurrentBuffers == true ? curBuffer : 
    this->curBuffer == 0 ? 1 : 0;

  _computeTotalEnergy[devType](this->devNum,
                               this->x[bufferToUse],
                               this->y[bufferToUse],
                               this->z[bufferToUse],
                               this->xDim * this->yDim * this->zDim,
                               &this->waveFunction[bufferToUse],
                               &this->potential[bufferToUse],
                               &this->kinetic[bufferToUse]);
}

void Grid::initializeFunctionTable()
{
  _initializeRuntime[kCPU] = GridCPUInit;
  _computeTotalEnergy[kCPU] = GridCPUComputeTotalEnergy;
  _deinitializeRuntime[kCPU] = GridCPUDeinit;

#ifdef USE_OPENCL
  _initializeRuntime[kOpenCL] = GridOpenCLInit;
  _computeTotalEnergy[kOpenCL] = GridOpenCLComputeTotalEnergy;
  _deinitializeRuntime[kOpenCL] = GridOpenCLDeinit;
#else
  _initializeRuntime[kOpenCL] = GridRuntimeNotImplemented;
  _computeTotalEnergy[kOpenCL] = GridRuntimeNotImplemented;
  _deinitializeRuntime[kOpenCL] = GridRuntimeNotImplemented;
#endif

  Grid::functionTableInit = true;
}

void Grid::initializeRuntime(GridDeviceType dev)
{
  if(Grid::functionTableInit == false)
  {
    Grid::initializeFunctionTable();
  }

  switch(dev)
  {
    case kCPU:
    case kOpenCL:
      _initializeRuntime[dev]();
      break;
    default:
      throw GridException("Grid::initializeRuntime():Unknown device type.");
  }
}

Grid::Grid(int xDim,
           int yDim,
           int zDim,
           float lengthScale, 
           GridDeviceType devType,
           int devNum,
           int* randomNumStream)
{
  int curParticle = 0;
  accepts = 0;
  totalIterations = 0;

  if(Grid::functionTableInit == false)
  {
    initializeFunctionTable();
  }

  this->randomNumStream = randomNumStream == NULL ?
    init_sprng(SPRNG_LFG, 0, 1, kRandomSeed, SPRNG_DEFAULT) :
    randomNumStream;

  this->lengthScale = lengthScale < 0 ? kStepSize : lengthScale;
  this->curBuffer = 0;
  this->devNum = devNum;
  this->devType = devType;
  this->xDim = xDim;
  this->yDim = yDim;
  this->zDim = zDim;

  try
  {
    this->x[0] = new float[xDim * yDim * zDim];
    this->y[0] = new float[xDim * yDim * zDim];
    this->z[0] = new float[xDim * yDim * zDim];
    this->x[1] = new float[xDim * yDim * zDim];
    this->y[1] = new float[xDim * yDim * zDim];
    this->z[1] = new float[xDim * yDim * zDim];
  }
  catch(std::exception)
  {
    throw GridException("Out of memory.");
  }

  for(int i = 0; i < xDim; i++)
  {
    for(int j = 0; j < yDim; j++)
    {
      for(int k = 0; k < zDim; k++)
      {
        this->x[0][curParticle] = i * 7.0;
        this->y[0][curParticle] = j * 7.0;
        this->z[0][curParticle++] = k * 7.0;
      }
    }
  }
  this->waveFunction[0] = 0.0f;
  this->waveFunction[1] = 0.0f;
}

void Grid::perturbGrid()
{
  int curParticle = 0;
  int propBuffer = curBuffer == 0 ? 1 : 0;

  for(int i = 0; i < xDim; i++)
  {
    for(int j = 0; j < yDim; j++)
    {
      for(int k = 0; k < zDim; k++)
      {
        this->x[propBuffer][curParticle] = 
          this->x[curBuffer][curParticle] +
          (sprng(this->randomNumStream) - 0.5) *
          this->lengthScale;
        this->y[propBuffer][curParticle] = 
          this->y[curBuffer][curParticle] +
          (sprng(this->randomNumStream) - 0.5) *
          this->lengthScale;
        this->z[propBuffer][curParticle] = 
          this->z[curBuffer][curParticle] +
          (sprng(this->randomNumStream) - 0.5) *
          this->lengthScale;
        curParticle++;
      }
    }
  }
}

void Grid::accept()
{
  curBuffer = curBuffer == 0 ? 1 : 0;
  accepts++;
  totalIterations++;
}

void Grid::reject()
{
  totalIterations++;
}

Grid::~Grid()
{
  delete[] x[0];
  delete[] x[1];
  delete[] y[0];
  delete[] y[1];
  delete[] z[0];
  delete[] z[1];
}

void Grid::printMetrics()
{
  double maxDistance = 0.0;
  double minDistance = 20e99;

  for(size_t i = 0; i < xDim * yDim * zDim; i++)
  {
    for(size_t j = 0; j < i; j++)
    {
      double dx = x[curBuffer][i] - x[curBuffer][j];
      double dy = y[curBuffer][i] - y[curBuffer][j];
      double dz = z[curBuffer][i] - z[curBuffer][j];

      double d = sqrt(dx*dx + dy*dy + dz*dz);

      maxDistance = d > maxDistance ? d : maxDistance;
      minDistance = d < minDistance ? d : minDistance;
    }
  }

  double acceptRatio = (double)accepts / (double)totalIterations;
  
  printf("max distance:%lE min distance:%lE accept ratio:%lE\n", 
         maxDistance, 
         minDistance,
         acceptRatio);
}
